/*
 *  net.c : Simple network layer with TFTP/ICMP builtin
 *
 *  Copyright (c) 2003, Intel Corporation (yu.tang@intel.com)
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License version 2 as
 *  published by the Free Software Foundation.
 */

#ifdef HAVE_CONFIG_H
# include <blob/config.h>
#endif

#include <blob/arch.h>
#include <blob/command.h>
#include <blob/serial.h>
#include <blob/util.h>
#include <blob/time.h>

#include <net.h>

#ifdef DEBUG
#define DBPRINT(args...) printf(args)
#else
#define DBPRINT(args...)
#endif

static struct mybuf in = {0} ,out1 = {0}, out2 = {0};
static unsigned char broadcast_mac_addr[6] = { 0xff, 0xff, 0xff, 0xff, 0xff, 0xff};
static unsigned char broadcast_ip_addr[4] = {0xff, 0xff, 0xff, 0xff};

//static unsigned char mac_addr[6] = {0x0, 0x0, 0x0, 0x0, 0xf, 0xe};
static unsigned char ip_addr[4] = {192, 168, 1, 101};

static unsigned char svr_mac_addr[6];
static unsigned char svr_ip_addr[4] = {192,168,1,100};
static unsigned short our_tftp_port = 0;

static char image_name[256];
static unsigned int image_addr;

static int last_block_num;
static unsigned short ip_id=0; 

static unsigned long tftp_timeout = 0;
static int tftp_state = 0;
static int arp_state = 0;

#define MIN(a,b) ((a)<(b) ? (a): (b))

static unsigned short in_chksum(unsigned char *p, int len);
static void arp_input();
static void ip_input();
static void icmp_input();
static void udp_input();
static void arp_request();
static void tftp_handler();
static void do_tftp();

/* 
 * Very simple buffer management.
 *
 * io = 0 : request for input buffer
 * io = 1 : request for output buffer
 *
 * set buffer length to 0 to free the buffer.
 *
 */
struct mybuf * bget(int io) 
{
	if( io == 1) {	
		if(out1.len == 0) return &out1;
		if(out2.len == 0) return &out2;
	}else return &in;

	printf("%s:can't get buffer\n",__FUNCTION__);
	return 0;
}

static unsigned short in_chksum(unsigned char *p, int len)
{
	unsigned long sum=0;
	
	while(len > 1) {
		sum += *((unsigned short*)p)++;
		if( sum & 0x80000000 ) 
			sum = (sum & 0xFFFF) + (sum >> 16);
		len -= 2;
	}

	if(len) 
		sum += (unsigned short) *(unsigned char*) p;

	while(sum >> 16) 
		sum = (sum & 0xFFFF) + (sum >> 16);

	return ~sum;
}



/* ARP zone */

static void arp_input()
{
	int op;
	struct etherhdr *eh, *ehout;
	struct ether_arp *ea, *eaout ;
	struct mybuf *out;

	out = bget(1);
	if(!out) return ;

	eh = (struct etherhdr*) (in.ether);
	ea = (struct ether_arp*)(eh + 1);

	ehout = (struct etherhdr*) (out->buf);
	eaout = (struct ether_arp*) (ehout + 1);

	/* sanity check */
	if ( ntohs(ea->ea_hdr.ar_hrd) != ARPHDR_ETHER ) {
		DBPRINT("check 1\n");
	       	return;
	}
	if ( ntohs(ea->ea_hdr.ar_pro) != ETHPROTO_IP ) {
		DBPRINT("check 2\n");
		return;
	}
	if ( ea->ea_hdr.ar_hln != 6) {
		DBPRINT("check 3\n");
		return;
	}

	if ( ea->ea_hdr.ar_pln != 4) {
		DBPRINT("check 4\n");
		return;
	}

	/* matching ip */

	op = ntohs(ea->ea_hdr.ar_op);

	switch(op) {
	case ARP_REQUEST:
		*ehout = *eh;
		*eaout = *ea;

		if ( !memcmp(ea->arp_spa, svr_ip_addr, 4) ) {
			memcpy(svr_mac_addr, ea->arp_sha, 6);
		}

		memcpy(ehout->eh_dhost, eh->eh_shost, 6);
		memcpy(ehout->eh_shost, eth_mac_addr, 6);

		eaout->ea_hdr.ar_op = htons(ARP_REPLY);

		memcpy(eaout->arp_tha, ea->arp_sha, 6);
		memcpy(eaout->arp_tpa, ea->arp_spa, 4);

		memcpy(eaout->arp_sha, eth_mac_addr, 6);
		memcpy(eaout->arp_spa, ip_addr, 4);

		out->len = sizeof(struct etherhdr) + sizeof(struct ether_arp);

		/* send it */
		eth_xmit(out);

		break;
	case ARP_REPLY:
		if( (!memcmp(eh->eh_dhost, eth_mac_addr, 6)) && ( !memcmp(ea->arp_spa, svr_ip_addr, 4) ) ) {
			arp_state = 1;
			memcpy(svr_mac_addr, ea->arp_sha, 6);
			do_tftp();
		}
		break;
	case ARP_REVREQUEST:
		break;
	case ARP_REVREPLY:
		break;
	}
}

static void arp_request()
{	
	struct etherhdr *eh;
	struct ether_arp *ea;	
	struct mybuf *out;

	out = bget(1);
	if(!out) return;

	eh = (struct etherhdr *) (out->buf);
	ea = (struct ether_arp *) (eh + 1 );

	memcpy(eh->eh_dhost, broadcast_mac_addr, 6);
	memcpy(eh->eh_shost, eth_mac_addr, 6);
	eh->eh_proto = htons( ETHPROTO_ARP);
	
	ea->ea_hdr.ar_hrd = htons(ARPHDR_ETHER);
	ea->ea_hdr.ar_pro = htons(ETHPROTO_IP);
	ea->ea_hdr.ar_hln = 6;
	ea->ea_hdr.ar_pln = 4;
	ea->ea_hdr.ar_op = htons(ARP_REQUEST);

	memcpy(ea->arp_sha, eth_mac_addr, 6);
	memcpy(ea->arp_spa, ip_addr, 4);
	//memcpy(ea->.tha, 6);
	memcpy(ea->arp_tpa, svr_ip_addr, 4);

	out->len = sizeof(struct etherhdr) + sizeof(struct ether_arp);

	eth_xmit(out);
}

/* IP zone */
static void ip_input()
{
	int hlen;
	struct iphdr *ip;

	ip = in.ip = (struct iphdr*) (in.ether+1);

	if (ip->ip_v != IPVERSION) {
		printf("%s:ip version not matacing\n", __FUNCTION__);
		return;
	}

	hlen =  ip->ip_hl << 2 ;

	if( hlen < sizeof(struct iphdr) ) {
		printf("%s:hlen < iphdr\n", __FUNCTION__);
		return;
	}

	ip->ip_sum = in_chksum((unsigned char*)ip, hlen);
       	if(ip->ip_sum) {
		printf("%s:chksum error\n",__FUNCTION__);
		return;
	}
	
	if(ntohs(ip->ip_len) < hlen) {
		printf("%s:ip_len < hlen: 0x%x, 0x%x\n",__FUNCTION__,ip->ip_len, hlen);
		return;
	}

	if( ntohs(ip->ip_off) & IP_OFFMASK ) {
		printf("%s:can't handler fragment\n", __FUNCTION__);
		return;
	}

	if( memcmp(&ip->ip_dst, ip_addr, 4) && memcmp(&ip->ip_dst,broadcast_ip_addr,4) ) {
		/* packet not for us, ingore */
		DBPRINT("noise, not for us\n");
		return;
	}

	switch(ip->ip_p) {
	case IPPROTO_ICMP:
		icmp_input();
		break;
	case IPPROTO_UDP:
		udp_input();
		break;
	default:
		DBPRINT("Unsupported IP packet\n");
		break;
	}
	return; 
}

/* ICMP zone */
static void icmp_input()
{
	struct in_addr t;
	struct icmphdr *icmp;
	struct mybuf *out;

	struct etherhdr *eh = in.ether;
	struct iphdr *ip = in.ip;
	int hlen = ip->ip_hl << 2;
	int icmplen = ntohs(ip->ip_len) - hlen;

	if( icmplen < ICMP_MINLEN ) {
		printf("%s: .. < ICMP_MINLEN\n", __FUNCTION__);
		return;
	}

	icmp = (struct icmphdr *) ((u8*)in.ip + hlen);

	switch(icmp->type) {
	case ICMP_ECHO:
		out = bget(1);
		if(!out) return;

		icmp->type = ICMP_ECHOREPLY;
		icmp->cksum = 0;
		icmp->cksum = in_chksum((unsigned char*)icmp, icmplen);
		memcpy(out->buf + sizeof(struct etherhdr) + hlen, 
				icmp, icmplen);

		memcpy(&t, &ip->ip_dst, sizeof(struct in_addr));
		memcpy(&ip->ip_dst, &ip->ip_src, sizeof(struct in_addr));
		memcpy(&ip->ip_src, &t, sizeof(struct in_addr)); 

		ip->ip_ttl = MAXTTL;
		ip->ip_hl = sizeof(struct iphdr) >> 2; /* discard the options */
		ip->ip_len = htons(icmplen + sizeof(struct iphdr) );
		ip->ip_sum = 0;
		ip->ip_sum = in_chksum((unsigned char*)ip, sizeof(struct iphdr));
		memcpy(out->buf + sizeof(struct etherhdr), ip, sizeof(struct iphdr));

		memcpy(eh->eh_dhost, eh->eh_shost, 6);
		memcpy(eh->eh_shost, eth_mac_addr, 6);
		memcpy(out->buf, eh, sizeof(struct etherhdr) );

		out->len = icmplen + sizeof(struct iphdr) + sizeof( struct etherhdr);
		eth_xmit(out);

		break;

	case ICMP_ECHOREPLY:
		break;

	case ICMP_UNREACH:
		break;
	}

	return;

}

/* TFTP zone */
static void udp_input()
{
	struct etherhdr *eh;
	struct iphdr *ip;
	struct udphdr *uh;


	eh = in.ether;
	ip = in.ip;
	uh = in.udp = (struct udphdr *) ( (u8*)ip + (ip->ip_hl << 2));

	/* check-sum */

	/* Not for us? */
	if( ntohs(uh->uh_dport) != our_tftp_port ) return;

	/* ok, tftp handler */
	tftp_handler();
}


#define	RRQ	01				/* read request */
#define	WRQ	02				/* write request */
#define	DATA	03				/* data packet */
#define	ACK	04				/* acknowledgement */
#define	ERROR	05				/* error code */


static void do_tftp()
{
	struct mybuf* out;
	char *p , *th ;
	char mode[]="octet";
	int pktlen = 0;
	struct etherhdr *eh;
	struct iphdr *ip;
	struct udphdr *uh;

	/* in progress... */
	if(tftp_state) return ;
	else tftp_state = 1;

	out = bget(1);
	if(!out) return;

	eh = (struct etherhdr*) (out->buf);
	ip = (struct iphdr*) (eh + 1);
	uh = (struct udphdr*) (ip + 1);

	printf("downloading %s to 0x%8x\n", image_name, image_addr);

	/* reset */
	last_block_num = 0;
	our_tftp_port =0;


	/* tftp header */
	p = th = (unsigned char*)(uh+ 1);
	*(unsigned short *) p = htons(RRQ);
	p += 2;

	strcpy(p, image_name);
	p += strlen(image_name) + 1;

	strcpy(p, mode);
	p += sizeof(mode);

	pktlen = p - th;

	/* udp header */
	our_tftp_port = 1024 + ( TimerGetTime() % 3072);
	uh->uh_dport = htons(69);
	uh->uh_sport = htons(our_tftp_port);
	uh->uh_sum = 0;
	uh->uh_len = htons(pktlen + sizeof(struct udphdr) );

	/* ip header */
	ip->ip_v = IPVERSION;
	ip->ip_tos = 0; 
	ip->ip_hl = sizeof(struct iphdr) >> 2;
	ip->ip_id = htons(ip_id);
	ip_id++;
	ip->ip_off = htons(0x4000); /* no frag */
	ip->ip_p = IPPROTO_UDP;
	ip->ip_len = htons( pktlen + sizeof(struct udphdr) + sizeof(struct iphdr));
	ip->ip_ttl = MAXTTL;
	memcpy(&ip->ip_dst, svr_ip_addr,sizeof(struct in_addr));
	memcpy(&ip->ip_src, ip_addr, sizeof(struct in_addr));
	ip->ip_sum = 0;
	ip->ip_sum = in_chksum((unsigned char*)ip, sizeof(struct iphdr));

	/* ether header */
	memcpy(eh->eh_shost, eth_mac_addr, ETH_ALEN);
	memcpy(eh->eh_dhost, svr_mac_addr, ETH_ALEN);
	eh->eh_proto = htons(ETHPROTO_IP);

	out->len = pktlen + sizeof(struct udphdr) + sizeof(struct iphdr)+ sizeof(struct etherhdr);

	/* xmit */
	eth_xmit(out); 
}

static void tftp_handler()
{
	unsigned char *p, *th;
	unsigned short opcode, blocknum, expected;
	int pktlen, datalen;
	struct etherhdr *eh, *ehout;
	struct iphdr *ip, *ipout;
	struct udphdr *uh, *uhout;

	struct mybuf* out;

	/* Not in tftp session */
	if(!tftp_state) return;

	/* reset timeout count */
	tftp_timeout = 0;
	
	out = bget(1);
	if(!out) return;
	
	eh = in.ether;
	ip = in.ip;
	uh = in.udp;

	ehout = (struct etherhdr*)(out->buf);
	ipout = (struct iphdr*) ( ehout + 1);
	uhout = (struct udphdr*) (ipout + 1);

	p = th = (unsigned char*)uh + sizeof(struct udphdr);

	pktlen = ntohs(uh->uh_len) - sizeof(struct udphdr) ;


	if( pktlen < 2) return;

	opcode = ntohs(*(unsigned short*)p);
	p += 2;
	switch(opcode) {
	case RRQ:
	case WRQ:
		/* we are client ignore */
		break;
	case DATA:
		datalen = pktlen - 4;
		blocknum = ntohs(*(unsigned short*)p);
		p += 2;

		expected = (last_block_num + 1) % (2<<16);

		if ( last_block_num == blocknum ) {
			DBPRINT("same block\n");
		} 
		else if ( blocknum == expected )  {
			last_block_num = blocknum;
			/* store data */
			memcpy((unsigned char*)image_addr, p, datalen);
			image_addr += datalen; 
		}
		else {
			DBPRINT("TFTP : Expected block %d, got block %d\n", expected , blocknum);
			/* ASK FOR IT */
			blocknum = last_block_num;
		}

		th = (unsigned char*) (uhout + 1);
		*(unsigned short*) th = htons(ACK);
		*((unsigned short*) th + 1) = htons(blocknum);

		uhout->uh_dport = uh->uh_sport;
		uhout->uh_sport = htons(our_tftp_port);
		uhout->uh_len = htons(4 + sizeof(struct udphdr));
		uhout->uh_sum = 0;

		ipout->ip_v = IPVERSION;
		ipout->ip_tos = 0; 
		ipout->ip_hl = sizeof(struct iphdr) >> 2;
		ipout->ip_id = htons(ip_id);
		ip_id ++;
		ipout->ip_off = htons(0x4000); /* no frag */
		ipout->ip_p = IPPROTO_UDP;
		ipout->ip_len = htons( 4 + sizeof(struct udphdr) + sizeof(struct iphdr));
		ipout->ip_ttl = MAXTTL;
		memcpy(&ipout->ip_dst, svr_ip_addr,sizeof(struct in_addr));
		memcpy(&ipout->ip_src, ip_addr, sizeof(struct in_addr));
		ipout->ip_sum = 0;
		ipout->ip_sum = in_chksum((unsigned char*)ipout, sizeof(struct iphdr));

		memcpy(ehout->eh_shost, eth_mac_addr, ETH_ALEN);
		memcpy(ehout->eh_dhost, svr_mac_addr, ETH_ALEN);
		ehout->eh_proto = htons(ETHPROTO_IP);

		out->len = sizeof(struct iphdr) + sizeof(struct etherhdr) + sizeof(struct udphdr) + 4;
		/* ack it */
		eth_xmit(out);

		if (datalen < 512) {
			/* reset tftp state */
			tftp_state++;
			printf("done\n");
		}
		else {
			//printf(".");
			//if( 0 == (blocknum % 80) ) printf("\n");
		}
		break;
	case ERROR:
		tftp_state ++;
		printf("tftp error, aborted\n");

		break;
	}
}

void net_reset()
{
	out1.len = 0;
	out2.len = 0;
	in.len = 0;
}

void net_rx() 
{
	struct etherhdr *eh;
	int proto;

	eh = (struct etherhdr*)in.buf;
	
	proto = ntohs(eh->eh_proto);
	/* IEEE 802 packet */
	if (proto < 1514) {
		memcpy((u8*)eh->eh_shost, (u8*)(eh+1), ETH_ALEN);
		memcpy((u8*)eh->eh_dhost, ((u8*)(eh+1)) - 6, ETH_ALEN);

		in.ether = (struct etherhdr*) (in.buf + 8);
	}
	else 
		in.ether = eh;

	proto = ntohs(eh->eh_proto);

	switch(proto) {
	case ETHPROTO_IP:
		ip_input();
		break;
	case ETHPROTO_ARP:
		arp_input() ;
		break;
	default:
		DBPRINT("Unsupported packet : 0x%x\n", proto);
		break;
	}

}

int cmd_tftp(int argc, char* argv[])
{
	int try = 3;
	unsigned long timeout = TICKS_PER_SECOND;

	if(argc < 3) {
		printf("usage:\n \ttftp file-name ram-addr\n");
		return 0;
	}

	strcpy(image_name, argv[1]);
	strtou32(argv[2], (u32*)&image_addr);

	/* reset count */
	arp_state = 0;
	tftp_state = 0;
	tftp_timeout = 0;

	while(!arp_state) {
		timeout = TICKS_PER_SECOND/10;
		arp_request();
		while(timeout && !arp_state) {
			eth_rx();
			timeout -- ;
		}

		if(arp_state) break;
		if( !(try--) ) {
			printf("TFTP : ARP timeout.\n");
			break;
		}
	}

	while(tftp_state < 2) {
		tftp_timeout ++ ;
		if(tftp_timeout > TICKS_PER_SECOND) {
			printf("TFTP : transfer timeout\n");
			break;
		}

		eth_rx();
	}

	/* reset count */
	arp_state = 0;
	tftp_state = 0;
	tftp_timeout = 0;

	return 0;
}

static char flashtftphelp[] = "tftp file-name ram-addr\n" ;
__commandlist( cmd_tftp, "tftp", flashtftphelp );


static int str2mac(unsigned char* str, unsigned char* mac)
{
	int i=0;
	unsigned char *p=str;
	unsigned char t[6];
	int addr;
	
	while(i < 6) {
		addr = 0;
		while( *p && (*p != ':') ) {
			if( (*p >= '0') && (*p <= '9') ) {
				addr = addr * 16 + (*p - '0');
				p++;
			}
			else if( (*p >= 'a') && (*p <= 'f') ) {
				addr = addr * 16 + (*p - 'a');
				p++;
			}
			else if ( (*p >= 'A') && (*p <= 'F') ) {
				addr = addr * 16 + (*p - 'A');
			       	p++;
			}
			else return -1; /* invalid */
		}

		if (addr < 255) t[i] = (addr&0xFF);
		else break;

		i++;

		if(*p) p++;
		else break;
	}

	if( i!=6 )  return -1;

	memcpy(mac, t, sizeof(t));

	return 0;
}

static int str2ip(unsigned char* str, unsigned char* ip)
{
	int i=0;
	unsigned char *p=str;
	unsigned char t[4];
	int addr;
	
	while(i < 4) {
		addr = 0;
		while( *p && (*p != '.') ) {
			if( (*p >= '0') && (*p <= '9') ) {
				addr = addr * 10 + (*p - '0');
			       	p++;
			}
			else return -1; /* invalid */
		}

		if (addr < 255) t[i] = (addr&0xFF);
		else break;

		i++;

		if(*p) p++;
		else break;

	}

	if( i!=4 )  return -1;

	memcpy(ip, t, sizeof(t));

	return 0;
}

static int cmd_ifconfig(int argc, char* argv[])
{
	if(argc == 1) {
		printf("Mac addr\t: %02x.%02x.%02x.%02x.%02x.%02x\n",
			eth_mac_addr[0],eth_mac_addr[1],eth_mac_addr[2],
			eth_mac_addr[3],eth_mac_addr[4],eth_mac_addr[5]);
		printf("Our IP addr\t: %d.%d.%d.%d\n", ip_addr[0], ip_addr[1], ip_addr[2], ip_addr[3]);
		printf("Server IP addr\t: %d.%d.%d.%d\n", svr_ip_addr[0],
				svr_ip_addr[1], svr_ip_addr[2], svr_ip_addr[3]);
		return 0;

	}else if(argc < 3) {
		printf("usage:\n \t ifconfig [ip|mac|server] addr\n");
		return 0;
	}

	if(!strncmp(argv[1], "ip", 6)) {
		return str2ip(argv[2],ip_addr);
	}

	if(!strncmp(argv[1],"mac", 6)) {
		return str2mac(argv[2],eth_mac_addr);
	}

	if(!strncmp(argv[1],"server", 6)) {
		return str2ip(argv[2],svr_ip_addr);

	}
	return -1;

}

static char flashifconfighelp[] = "ifconfig [ip|mac|server] addr\n" ;
__commandlist( cmd_ifconfig, "ifconfig", flashifconfighelp );

